import torch
from torch.utils.data import DataLoader, Dataset


class VOCDataset(Dataset):
    def __init__(self, file_path):
        super().__init__()
        data = torch.load(file_path, weights_only=True)
        self.features = data['features']
        self.labels = data['labels']

    def __getitem__(self, index):
        return self.features[index], self.labels[index]

    def __len__(self):
        return len(self.labels)


def generate_loader(path: str, batch_size: int, shuffle: bool):
    dataset = VOCDataset(path)
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
